{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp run_bert_multitask\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Bert Multitask Learning\n",
    "\n",
    "Train, eval and predict api for bert multitask learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import json\n",
    "import os\n",
    "import time\n",
    "from shutil import copytree, ignore_patterns, rmtree\n",
    "from typing import Callable, Dict, List, Tuple, Union\n",
    "\n",
    "from sklearn.preprocessing import MultiLabelBinarizer\n",
    "from transformers import PreTrainedTokenizer\n",
    "\n",
    "import tensorflow as tf\n",
    "from loguru import logger\n",
    "from m3tl.input_fn import predict_input_fn, train_eval_input_fn\n",
    "from m3tl.model_fn import BertMultiTask\n",
    "from m3tl.params import Params\n",
    "from m3tl.special_tokens import EVAL, PREDICT\n",
    "from m3tl.utils import (compress_tf_warnings, get_or_make_label_encoder,\n",
    "                        infer_shape_and_type_from_dict, set_phase, LabelEncoder, get_is_pyspark)\n",
    "from tensorflow.python.framework.errors_impl import \\\n",
    "    NotFoundError as TFNotFoundError\n",
    "\n",
    "compress_tf_warnings()\n",
    "# Fix duplicate log\n",
    "# LOGGER = tf.get_logger()\n",
    "# LOGGER.propagate = False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:22:35.631 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:22:35.632 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:22:35.632 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:22:35.633 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:22:35.633 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_regression, problem type: regression\n",
      "2021-06-17 13:22:35.634 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "2021-06-17 13:22:35.635 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:22:35.636 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem fake_contrastive_learning, problem type: contrastive_learning\n",
      "2021-06-17 13:22:35.637 | WARNING  | m3tl.base_params:assign_problem:634 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-17 13:22:35.638 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-17 13:22:41.256 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "import tensorflow as tf\n",
    "\n",
    "from m3tl.predefined_problems import *\n",
    "\n",
    "from m3tl import Params\n",
    "import os\n",
    "from m3tl import predict_input_fn\n",
    "params = Params()\n",
    "params.shuffle_buffer = 1000\n",
    "\n",
    "# configure transformers\n",
    "params.transformer_tokenizer_loading = 'BertTokenizer'\n",
    "params.transformer_model_loading = 'AlbertForMaskedLM'\n",
    "params.transformer_config_loading = 'AlbertConfig'\n",
    "params.transformer_model_name = 'voidful/albert_chinese_tiny'\n",
    "params.transformer_config_name = 'voidful/albert_chinese_tiny'\n",
    "params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'\n",
    "\n",
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "\n",
    "tb = TestBase()\n",
    "params = tb.params\n",
    "\n",
    "# hide\n",
    "problem = 'weibo_fake_ner&weibo_fake_cls|weibo_fake_multi_cls|weibo_masklm|weibo_premask_mlm'\n",
    "problem_type_dict = {\n",
    "    'weibo_fake_ner': 'seq_tag',\n",
    "    'weibo_cws': 'seq_tag',\n",
    "    'weibo_fake_multi_cls': 'multi_cls',\n",
    "    'weibo_fake_cls': 'cls',\n",
    "    'weibo_masklm': 'masklm',\n",
    "    'weibo_pretrain': 'pretrain',\n",
    "    'weibo_premask_mlm': 'premask_mlm'\n",
    "}\n",
    "\n",
    "processing_fn_dict = {\n",
    "    'weibo_fake_ner': get_weibo_fake_ner_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "    'weibo_cws': get_weibo_cws_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "    'weibo_fake_cls': get_weibo_fake_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "    'weibo_fake_multi_cls': get_weibo_fake_multi_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "    'weibo_masklm': get_weibo_masklm(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "    'weibo_pretrain': get_weibo_pretrain_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "    'weibo_premask_mlm': get_weibo_premask_mlm(file_path='/data/m3tl/data/ner/weiboNER*')\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def create_keras_model(\n",
    "        mirrored_strategy: tf.distribute.MirroredStrategy,\n",
    "        params: Params,\n",
    "        mode='train',\n",
    "        inputs_to_build_model=None,\n",
    "        model=None,\n",
    "        run_eagerly=False):\n",
    "\n",
    "    def _get_model_wrapper(params, mode, inputs_to_build_model, model):\n",
    "        # Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow\n",
    "        # uses hvd.DistributedOptimizer() to compute gradients.\n",
    "        experimental_run_tf_function = not params.use_horovod\n",
    "        if model is None:\n",
    "            model = BertMultiTask(params)\n",
    "            # model.run_eagerly = True\n",
    "        set_phase(PREDICT)\n",
    "        if mode == 'resume':\n",
    "            model.compile(run_eagerly=run_eagerly,\n",
    "                          experimental_run_tf_function=experimental_run_tf_function)\n",
    "            # build training graph\n",
    "            # model.train_step(inputs_to_build_model)\n",
    "            \n",
    "            _ = model(inputs_to_build_model)\n",
    "            # load ALL vars including optimizers' states\n",
    "            try:\n",
    "                model.load_weights(os.path.join(\n",
    "                    params.ckpt_dir, 'model'), skip_mismatch=False)\n",
    "            except TFNotFoundError:\n",
    "                logger.warning('Not resuming since no mathcing ckpt found')\n",
    "        elif mode == 'transfer':\n",
    "            # build graph without optimizers' states\n",
    "            # calling compile again should reset optimizers' states but we're playing safe here\n",
    "            _ = model(inputs_to_build_model)\n",
    "            # load weights without loading optimizers' vars\n",
    "            model.load_weights(os.path.join(params.init_checkpoint, 'model'))\n",
    "            # compile again\n",
    "            model.compile(run_eagerly=run_eagerly,\n",
    "                          experimental_run_tf_function=experimental_run_tf_function)\n",
    "        elif mode == 'predict':\n",
    "            _ = model(inputs_to_build_model)\n",
    "            # load weights without loading optimizers' vars\n",
    "            model.load_weights(os.path.join(params.ckpt_dir, 'model'))\n",
    "        elif mode == 'eval':\n",
    "            _ = model(inputs_to_build_model)\n",
    "            # load weights without loading optimizers' vars\n",
    "            model.load_weights(os.path.join(params.ckpt_dir, 'model'))\n",
    "            model.compile(run_eagerly=run_eagerly,\n",
    "                          experimental_run_tf_function=experimental_run_tf_function)\n",
    "        else:\n",
    "            model.compile(run_eagerly=run_eagerly,\n",
    "                          experimental_run_tf_function=experimental_run_tf_function)\n",
    "\n",
    "        return model\n",
    "    if mirrored_strategy is not None:\n",
    "        with mirrored_strategy.scope():\n",
    "            model = _get_model_wrapper(\n",
    "                params, mode, inputs_to_build_model, model)\n",
    "    else:\n",
    "        model = _get_model_wrapper(params, mode, inputs_to_build_model, model)\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Init model in various mode\n",
    "\n",
    "`train`: model will be loaded from huggingface\n",
    "`resume`: model will be loaded from params.ckpt_dir, if params.ckpt_dir dose not contain valid checkpoint, then load from huggingface\n",
    "`transfer`: model will be loaded from params.init_checkpoint, the correspongding path should contain checkpoints saved using m3tl\n",
    "`predict`: model will be loaded from params.ckpt_dir except optimizers' states\n",
    "`eval`: model will be loaded from params.ckpt_dir except optimizers' states, model will be compiled\n",
    "\n",
    "Args:\n",
    "- mirrored_strategy (tf.distribute.MirroredStrategy): mirrored strategy\n",
    "- params (Params): params\n",
    "- mode (str, optional): Mode, see above explaination. Defaults to 'train'.\n",
    "- inputs_to_build_model (Dict, optional): A batch of data. Defaults to None.\n",
    "- model (Model, optional): Keras model. Defaults to None.\n",
    "\n",
    "Returns:\n",
    "- model: loaded model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def _has_callbacks(callbacks: List[tf.keras.callbacks.Callback], check_callback: tf.keras.callbacks.Callback) -> bool:\n",
    "    for callback in callbacks:\n",
    "        if isinstance(callback, check_callback):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "\n",
    "def _train_bert_multitask_keras_model(train_dataset: tf.data.Dataset,\n",
    "                                      eval_dataset: tf.data.Dataset,\n",
    "                                      model: tf.keras.Model,\n",
    "                                      params: Params,\n",
    "                                      mirrored_strategy: tf.distribute.MirroredStrategy = None,\n",
    "                                      callbacks: List[tf.keras.callbacks.Callback] = None,\n",
    "                                      verbose=1):\n",
    "\n",
    "    all_callbacks = params.gather_mtl_callbacks()\n",
    "\n",
    "    if callbacks is not None:\n",
    "        all_callbacks += callbacks\n",
    "\n",
    "    # if callbacks is not passed or callbacks dose not contain\n",
    "    # ModelCheckpoint and TensorBoard callbacks, we add the default ones\n",
    "\n",
    "    # can't save whole model with model subclassing api due to tf bug\n",
    "    # see: https://github.com/tensorflow/tensorflow/issues/42741\n",
    "    # https://github.com/tensorflow/tensorflow/issues/40366\n",
    "    model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "        filepath=os.path.join(params.ckpt_dir, 'model'),\n",
    "        save_weights_only=True,\n",
    "        monitor='val_mean_acc',\n",
    "        mode='auto',\n",
    "        save_best_only=False)\n",
    "\n",
    "    tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
    "        log_dir=params.ckpt_dir)\n",
    "\n",
    "    has_model_checkpoint_callback = _has_callbacks(\n",
    "        all_callbacks, tf.keras.callbacks.ModelCheckpoint)\n",
    "\n",
    "    # horovod callbacks\n",
    "    if params.use_horovod:\n",
    "        import horovod.tensorflow.keras as hvd\n",
    "        # when using horovod as dist backend, only save model in process 0\n",
    "        if not has_model_checkpoint_callback:\n",
    "            if hvd.rank() == 0:\n",
    "                all_callbacks.append(model_checkpoint_callback)\n",
    "        all_callbacks = [\n",
    "            hvd.callbacks.BroadcastGlobalVariablesCallback(0),\n",
    "            hvd.callbacks.MetricAverageCallback()\n",
    "        ] + all_callbacks\n",
    "        # Horovod: write logs on worker 0.\n",
    "        verbose = verbose if hvd.rank() == 0 else 0\n",
    "    elif not has_model_checkpoint_callback:\n",
    "        all_callbacks.append(model_checkpoint_callback)\n",
    "\n",
    "    validation_steps = params.get('validation_steps', 1000)\n",
    "\n",
    "    if mirrored_strategy is not None:\n",
    "        with mirrored_strategy.scope():\n",
    "            model.fit(\n",
    "                x=train_dataset,\n",
    "                validation_data=eval_dataset,\n",
    "                epochs=params.train_epoch,\n",
    "                callbacks=all_callbacks,\n",
    "                steps_per_epoch=params.train_steps_per_epoch,\n",
    "                verbose=verbose,\n",
    "                validation_steps=validation_steps\n",
    "            )\n",
    "    else:\n",
    "        model.fit(\n",
    "            x=train_dataset,\n",
    "            validation_data=eval_dataset,\n",
    "            epochs=params.train_epoch,\n",
    "            callbacks=all_callbacks,\n",
    "            steps_per_epoch=params.train_steps_per_epoch,\n",
    "            verbose=verbose,\n",
    "            validation_steps=validation_steps\n",
    "        )\n",
    "    model.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_params_ready(problem, num_gpus, model_dir, params, problem_type_dict, processing_fn_dict, mode='train', json_path='') -> Params:\n",
    "    if params is None:\n",
    "        params = Params()\n",
    "    if not os.path.exists('models'):\n",
    "        os.mkdir('models')\n",
    "    if model_dir:\n",
    "        base_dir, dir_name = os.path.split(model_dir)\n",
    "    else:\n",
    "        base_dir, dir_name = None, None\n",
    "    # add new problem to params if problem_type_dict and processing_fn_dict provided\n",
    "    if problem_type_dict:\n",
    "        params.register_multiple_problems(\n",
    "            problem_type_dict=problem_type_dict, processing_fn_dict=processing_fn_dict)\n",
    "    \n",
    "    if not params.problem_assigned and not problem:\n",
    "        raise ValueError('neither params problem assigned nor problem provided.')\n",
    "\n",
    "    if mode == 'train':\n",
    "        if problem:\n",
    "            params.assign_problem(problem, model_dir=model_dir)\n",
    "        params.to_json()\n",
    "    else:\n",
    "        params.from_json(json_path)\n",
    "        if problem:\n",
    "            params.assign_problem(problem, model_dir=model_dir, predicting=True)\n",
    "\n",
    "    return params\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# @logger.catch\n",
    "def train_bert_multitask(\n",
    "        problem='weibo_ner',\n",
    "        num_gpus=1,\n",
    "        num_epochs=10,\n",
    "        model_dir='',\n",
    "        params: Params = None,\n",
    "        problem_type_dict: Dict[str, str] = None,\n",
    "        processing_fn_dict: Dict[str, Callable] = None,\n",
    "        model: tf.keras.Model = None,\n",
    "        create_tf_record_only=False,\n",
    "        steps_per_epoch: int = None,\n",
    "        warmup_ratio=0.1,\n",
    "        continue_training=False,\n",
    "        mirrored_strategy: tf.distribute.MirroredStrategy = None,\n",
    "        run_eagerly=False,\n",
    "        callbacks: List[tf.keras.callbacks.Callback] = None,\n",
    "        verbose=1) -> tf.keras.Model:\n",
    "    \"\"\"\n",
    "    Train Multi-task Bert model\n",
    "\n",
    "    Keyword Arguments:\n",
    "    - problem (str, optional) -- Problems to train. Defaults to 'weibo_ner'\n",
    "    - num_gpus (int, optional) -- Number of GPU to use. Defaults to 1\n",
    "    - num_epochs (int, optional) -- Number of epochs to train. Defaults to 10\n",
    "    - model_dir (str, optional) -- model dir. Defaults to ''\n",
    "    - params (Params, optional) -- Params to define training and models. Defaults to None\n",
    "    - problem_type_dict (dict, optional) -- Key: problem name, value: problem type. Defaults to None\n",
    "    - processing_fn_dict (dict, optional) -- Key: problem name, value: problem data preprocessing fn. Defaults to None\n",
    "    - model (tf.keras.Model, optional): if not provided, it will be created using `create_keras_model`. Defaults to None.\n",
    "    - create_tf_record_only (bool, optional): if `True`, the function will only create TFRecord without training model. Defaults to False.\n",
    "    - steps_per_epoch (int, optional): steps per epochs, if not provided, train datset will be looped once to calculate steps per epoch. Defaults to None.\n",
    "    - warmup_ratio (float, optional): lr warmup ratio. Defaults to 0.1.\n",
    "    - continue_training (bool, optional): whether to resume training from `model_dir`. Defaults to False.\n",
    "    - mirrored_strategy (MirroredStrategy, optional): Tensorflow MirroredStrategy. Defaults to None.\n",
    "    - run_eagerly (bool, optional): Whether to run model eagerly. Defaults to False.\n",
    "    - callbacks (list, optional): list of callbacks to add during training. If None, ModelCheckpoint will be added.\n",
    "    - verbose (int, optional): 0, 1, or 2. Verbosity mode. 0 = silent, 1 = progress bar, 2 = one line per epoch. Note that the progress bar is not particularly useful when logged to a file, so verbose=2 is recommended when not running interactively (eg, in a production environment). \n",
    "    \"\"\"\n",
    "    if params is None:\n",
    "        params = Params()\n",
    "    if params.use_horovod:\n",
    "        import horovod.tensorflow.keras as hvd\n",
    "        hvd.init()\n",
    "        gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "        for gpu in gpus:\n",
    "            tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        if gpus:\n",
    "            tf.config.experimental.set_visible_devices(\n",
    "                gpus[hvd.local_rank()], 'GPU')\n",
    "\n",
    "    params = get_params_ready(problem, num_gpus, model_dir,\n",
    "                              params, problem_type_dict, processing_fn_dict)\n",
    "    params.train_epoch = num_epochs\n",
    "\n",
    "    train_dataset = train_eval_input_fn(params)\n",
    "    eval_dataset = train_eval_input_fn(params, mode=EVAL)\n",
    "    if create_tf_record_only:\n",
    "        return\n",
    "\n",
    "    if get_is_pyspark():\n",
    "        raise NotImplementedError(\n",
    "            'Pyspark only support creating TFRecord. Please set create_tf_record_only as True when pyspark is enabled.')\n",
    "\n",
    "    # get train_steps and update params\n",
    "    if steps_per_epoch is not None:\n",
    "        train_steps = steps_per_epoch\n",
    "    else:\n",
    "        train_steps = 0\n",
    "        for _ in train_dataset:\n",
    "            train_steps += 1\n",
    "    params.update_train_steps(train_steps, warmup_ratio=warmup_ratio)\n",
    "\n",
    "    train_dataset = train_eval_input_fn(params)\n",
    "    train_dataset = train_dataset.repeat()\n",
    "\n",
    "    one_batch = next(train_dataset.as_numpy_iterator())\n",
    "\n",
    "    if mirrored_strategy is None:\n",
    "        mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "    elif mirrored_strategy is False:\n",
    "        mirrored_strategy = None\n",
    "\n",
    "    if num_gpus > 1 and mirrored_strategy is not False:\n",
    "        train_dataset = mirrored_strategy.experimental_distribute_dataset(\n",
    "            train_dataset)\n",
    "        eval_dataset = mirrored_strategy.experimental_distribute_dataset(\n",
    "            eval_dataset)\n",
    "\n",
    "    # restore priority: self > transfer > huggingface\n",
    "    if continue_training and tf.train.latest_checkpoint(params.ckpt_dir):\n",
    "        mode = 'resume'\n",
    "    elif tf.train.latest_checkpoint(params.init_checkpoint):\n",
    "        mode = 'transfer'\n",
    "    else:\n",
    "        mode = 'train'\n",
    "\n",
    "    if model is None:\n",
    "        model = create_keras_model(\n",
    "            mirrored_strategy=mirrored_strategy, params=params,\n",
    "            mode=mode, inputs_to_build_model=one_batch,\n",
    "            run_eagerly=run_eagerly)\n",
    "\n",
    "    _train_bert_multitask_keras_model(\n",
    "        train_dataset=train_dataset,\n",
    "        eval_dataset=eval_dataset,\n",
    "        model=model,\n",
    "        params=params,\n",
    "        mirrored_strategy=mirrored_strategy,\n",
    "        callbacks=callbacks,\n",
    "        verbose=verbose\n",
    "    )\n",
    "    params.to_json()\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train Multi-task Bert model\n",
    "\n",
    "Keyword Arguments:\n",
    "- problem (str, optional) -- Problems to train. Defaults to 'weibo_ner'\n",
    "- num_gpus (int, optional) -- Number of GPU to use. Defaults to 1\n",
    "- num_epochs (int, optional) -- Number of epochs to train. Defaults to 10\n",
    "- model_dir (str, optional) -- model dir. Defaults to ''\n",
    "- params (Params, optional) -- Params to define training and models. Defaults to None\n",
    "- problem_type_dict (dict, optional) -- Key: problem name, value: problem type. Defaults to None\n",
    "- processing_fn_dict (dict, optional) -- Key: problem name, value: problem data preprocessing fn. Defaults to None\n",
    "- model (tf.keras.Model, optional): if not provided, it will be created using `create_keras_model`. Defaults to None.\n",
    "- create_tf_record_only (bool, optional): if `True`, the function will only create TFRecord without training model. Defaults to False.\n",
    "- steps_per_epoch (int, optional): steps per epochs, if not provided, train datset will be looped once to calculate steps per epoch. Defaults to None.\n",
    "- warmup_ratio (float, optional): lr warmup ratio. Defaults to 0.1.\n",
    "- continue_training (bool, optional): whether to resume training from `model_dir`. Defaults to False.\n",
    "- mirrored_strategy (MirroredStrategy, optional): Tensorflow MirroredStrategy. Defaults to None.\n",
    "- run_eagerly (bool, optional): Whether to run model eagerly. Defaults to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:22:41.927 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:22:41.928 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:22:41.928 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:22:41.929 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:22:41.929 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:22:41.930 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:22:41.930 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:22:41.931 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-17 13:22:42.005 | WARNING  | m3tl.read_write_tfrecord:chain_processed_data:250 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-17 13:22:42.012 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_fake_cls_weibo_fake_ner/train_00000.tfrecord\n",
      "2021-06-17 13:22:42.043 | WARNING  | m3tl.read_write_tfrecord:chain_processed_data:250 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-17 13:22:42.050 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_fake_cls_weibo_fake_ner/eval_00000.tfrecord\n",
      "2021-06-17 13:22:42.075 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_fake_multi_cls/train_00000.tfrecord\n",
      "2021-06-17 13:22:42.101 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_fake_multi_cls/eval_00000.tfrecord\n",
      "2021-06-17 13:22:42.176 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_masklm/train_00000.tfrecord\n",
      "2021-06-17 13:22:42.224 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_masklm/eval_00000.tfrecord\n",
      "2021-06-17 13:22:42.288 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_premask_mlm/train_00000.tfrecord\n",
      "2021-06-17 13:22:42.352 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmp4ngci4nl/weibo_premask_mlm/eval_00000.tfrecord\n",
      "2021-06-17 13:22:43.501 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:22:43.501 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:22:44.126 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:22:44.127 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:22:44.190 | CRITICAL | m3tl.base_params:update_train_steps:454 - Updating train_steps to 1\n",
      "2021-06-17 13:22:44.874 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:22:44.875 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.dense.bias', 'predictions.decoder.weight', 'predictions.LayerNorm.weight', 'predictions.bias', 'predictions.dense.weight', 'predictions.LayerNorm.bias', 'predictions.decoder.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-17 13:22:49.859 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:22:49.951 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:22:49.983 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-17 13:22:49.991 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-17 13:22:49.991 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-17 13:22:49.992 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "2021-06-17 13:22:50.941 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7fd82e8549f0>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - ETA: 0s - mean_acc: 3.6298 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1250 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.8964 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.9989 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.8322 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0389 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 10.0175"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:00.005 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:23:00.947 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1000 batches). You may need to use the repeat() function when building your dataset.\n",
      "1/1 [==============================] - 12s 12s/step - mean_acc: 3.6298 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1250 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.8964 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.9989 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.8322 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0389 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 10.0175 - val_loss: 14.9478 - val_mean_acc: 0.3333 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 1.0952 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.1085 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 3.0644 - val_BertMultiTaskTop/weibo_masklm/losses/0: 0.0000e+00\n",
      "Model: \"BertMultiTask\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "BertMultiTaskBody (BertMulti multiple                  4082696   \n",
      "_________________________________________________________________\n",
      "basic_mtl (BasicMTL)         multiple                  0         \n",
      "_________________________________________________________________\n",
      "BertMultiTaskTop (BertMultiT multiple                  13229575  \n",
      "_________________________________________________________________\n",
      "sum_loss_combination (SumLos multiple                  0         \n",
      "=================================================================\n",
      "Total params: 17,312,273\n",
      "Trainable params: 17,312,267\n",
      "Non-trainable params: 6\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "\n",
    "params.use_horovod = False\n",
    "\n",
    "model = train_bert_multitask(\n",
    "    problem=problem,\n",
    "    num_epochs=1,\n",
    "    params=params,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict,\n",
    "    steps_per_epoch=1,\n",
    "    continue_training=True,\n",
    "    mirrored_strategy=False,\n",
    "    run_eagerly=True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Slow train test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:02.297 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:23:02.297 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:23:02.298 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:23:02.299 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:23:02.299 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:23:02.300 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:23:02.300 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:23:02.301 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-17 13:23:02.922 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:02.922 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:23:03.734 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:03.734 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:23:03.800 | CRITICAL | m3tl.base_params:update_train_steps:454 - Updating train_steps to 1\n",
      "2021-06-17 13:23:04.386 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:04.387 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.dense.bias', 'predictions.decoder.weight', 'predictions.LayerNorm.weight', 'predictions.bias', 'predictions.dense.weight', 'predictions.LayerNorm.bias', 'predictions.decoder.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-17 13:23:09.611 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:23:09.704 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:23:09.735 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-17 13:23:09.742 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-17 13:23:09.743 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-17 13:23:09.743 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "2021-06-17 13:23:10.016 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: AutoGraph could not transform <bound method BertMultiTaskBody.call of <m3tl.model_fn.BertMultiTaskBody object at 0x7fd723ddc990>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: invalid value for \"node\": expected \"ast.AST\", got \"<class 'NoneType'>\"; to visit lists of nodes, use \"visit_block\" instead\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:23:17.066 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - ETA: 0s - mean_acc: 3.3779 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.4107 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.0287 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.4424 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.4780 - BertMultiTaskTop/weibo_masklm/losses/0: 9.9480 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.8374"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:23.802 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1000 batches). You may need to use the repeat() function when building your dataset.\n",
      "1/1 [==============================] - 17s 17s/step - mean_acc: 3.3779 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.4107 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.0287 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.4424 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.4780 - BertMultiTaskTop/weibo_masklm/losses/0: 9.9480 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.8374 - val_loss: 23.0261 - val_mean_acc: 0.4782 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.4286 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 2.1281 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.4307 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 1.2359 - val_BertMultiTaskTop/weibo_masklm/losses/0: 9.9194 - val_BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.8742\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:27.046 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:23:27.047 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:23:27.047 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:23:27.048 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:23:27.048 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:23:27.049 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:23:27.049 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:23:27.050 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"BertMultiTask\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "BertMultiTaskBody (BertMulti multiple                  4082696   \n",
      "_________________________________________________________________\n",
      "basic_mtl_1 (BasicMTL)       multiple                  0         \n",
      "_________________________________________________________________\n",
      "BertMultiTaskTop (BertMultiT multiple                  13229575  \n",
      "_________________________________________________________________\n",
      "sum_loss_combination_1 (SumL multiple                  0         \n",
      "=================================================================\n",
      "Total params: 17,312,273\n",
      "Trainable params: 17,312,267\n",
      "Non-trainable params: 6\n",
      "_________________________________________________________________\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:27.744 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:27.745 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:23:28.442 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:28.443 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:23:28.812 | CRITICAL | m3tl.base_params:update_train_steps:454 - Updating train_steps to 1\n",
      "2021-06-17 13:23:29.380 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:29.381 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.dense.bias', 'predictions.decoder.weight', 'predictions.LayerNorm.weight', 'predictions.bias', 'predictions.dense.weight', 'predictions.LayerNorm.bias', 'predictions.decoder.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-17 13:23:34.384 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:23:34.475 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:23:34.506 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-17 13:23:34.513 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-17 13:23:34.514 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-17 13:23:34.514 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:23:37.120 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:5 out of the last 5 calls to <function MultiModalBertModel.call at 0x7fd6e4e224d0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "1/1 [==============================] - ETA: 0s - mean_acc: 3.5559 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.0536 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.6014 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.0033 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.6906 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0329 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 10.0099"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:40.803 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:6 out of the last 6 calls to <function MultiModalBertModel.call at 0x7fd6e4e224d0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:41.726 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:7 out of the last 7 calls to <function MultiModalBertModel.call at 0x7fd6e4e224d0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1000 batches). You may need to use the repeat() function when building your dataset.\n",
      "1/1 [==============================] - 6s 6s/step - mean_acc: 3.5559 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.0536 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.6014 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.0033 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.6906 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0329 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 10.0099 - val_loss: 22.4883 - val_mean_acc: 0.2500 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_masklm/losses/0: 10.0292 - val_BertMultiTaskTop/weibo_premask_mlm/losses/0: 10.0646\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:43.008 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:23:43.008 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:23:43.009 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:23:43.009 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:23:43.010 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:23:43.010 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:23:43.011 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:23:43.011 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"BertMultiTask\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "BertMultiTaskBody (BertMulti multiple                  4082696   \n",
      "_________________________________________________________________\n",
      "basic_mtl_2 (BasicMTL)       multiple                  0         \n",
      "_________________________________________________________________\n",
      "BertMultiTaskTop (BertMultiT multiple                  13229575  \n",
      "_________________________________________________________________\n",
      "sum_loss_combination_2 (SumL multiple                  0         \n",
      "=================================================================\n",
      "Total params: 17,312,273\n",
      "Trainable params: 17,312,267\n",
      "Non-trainable params: 6\n",
      "_________________________________________________________________\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:43.644 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:43.645 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:23:44.289 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:44.290 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:23:45.277 | CRITICAL | m3tl.base_params:update_train_steps:454 - Updating train_steps to 2\n",
      "2021-06-17 13:23:46.258 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:23:46.259 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.dense.bias', 'predictions.decoder.weight', 'predictions.LayerNorm.weight', 'predictions.bias', 'predictions.dense.weight', 'predictions.LayerNorm.bias', 'predictions.decoder.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-17 13:23:51.472 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:23:51.570 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:23:51.603 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-17 13:23:51.612 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-17 13:23:51.612 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 2\n",
      "2021-06-17 13:23:51.613 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:8 out of the last 8 calls to <function MultiModalBertModel.call at 0x7fd6c0a0e5f0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:23:53.694 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:23:59.127 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - ETA: 0s - mean_acc: 3.3677 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.4085 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.0414 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.4213 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0993 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9891 "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:07.694 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1000 batches). You may need to use the repeat() function when building your dataset.\n",
      "2/2 [==============================] - 18s 5s/step - mean_acc: 3.3677 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.2724 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.0332 - BertMultiTaskTop/weibo_fake_ner/losses/0: 0.9476 - BertMultiTaskTop/weibo_masklm/losses/0: 10.1124 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9931 - val_loss: 22.4967 - val_mean_acc: 0.2500 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_masklm/losses/0: 10.0470 - val_BertMultiTaskTop/weibo_premask_mlm/losses/0: 10.0644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:12.062 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:24:12.063 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:24:12.063 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:24:12.064 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:24:12.064 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:24:12.064 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:24:12.065 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:24:12.065 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"BertMultiTask\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "BertMultiTaskBody (BertMulti multiple                  4082696   \n",
      "_________________________________________________________________\n",
      "basic_mtl_3 (BasicMTL)       multiple                  0         \n",
      "_________________________________________________________________\n",
      "BertMultiTaskTop (BertMultiT multiple                  13229575  \n",
      "_________________________________________________________________\n",
      "sum_loss_combination_3 (SumL multiple                  0         \n",
      "=================================================================\n",
      "Total params: 17,312,273\n",
      "Trainable params: 17,312,267\n",
      "Non-trainable params: 6\n",
      "_________________________________________________________________\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:12.694 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:24:12.695 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:24:13.327 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:24:13.328 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:24:13.390 | CRITICAL | m3tl.base_params:update_train_steps:454 - Updating train_steps to 1\n",
      "2021-06-17 13:24:13.948 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:24:13.949 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.dense.bias', 'predictions.decoder.weight', 'predictions.LayerNorm.weight', 'predictions.bias', 'predictions.dense.weight', 'predictions.LayerNorm.bias', 'predictions.decoder.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-17 13:24:18.935 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:24:19.061 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:24:19.092 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-17 13:24:19.099 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-17 13:24:19.099 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-17 13:24:19.100 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "2021-06-17 13:24:19.362 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:24:24.524 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - ETA: 0s - mean_acc: 3.5475 - weibo_fake_cls_acc: 0.3333 - weibo_fake_ner_acc: 0.0952 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.7740 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.7841 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.8613 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0436 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9406"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:31.396 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1000 batches). You may need to use the repeat() function when building your dataset.\n",
      "1/1 [==============================] - 15s 15s/step - mean_acc: 3.5475 - weibo_fake_cls_acc: 0.3333 - weibo_fake_ner_acc: 0.0952 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.7740 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.7841 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.8613 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0436 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9406 - val_loss: 22.2138 - val_mean_acc: 0.2500 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.7190 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_masklm/losses/0: 10.0376 - val_BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9525\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:34.681 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:24:34.682 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:24:34.683 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:24:34.684 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:24:34.684 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:24:34.685 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:24:34.685 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"BertMultiTask\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "BertMultiTaskBody (BertMulti multiple                  4082696   \n",
      "_________________________________________________________________\n",
      "basic_mtl_4 (BasicMTL)       multiple                  0         \n",
      "_________________________________________________________________\n",
      "BertMultiTaskTop (BertMultiT multiple                  13229575  \n",
      "_________________________________________________________________\n",
      "sum_loss_combination_4 (SumL multiple                  0         \n",
      "=================================================================\n",
      "Total params: 17,312,273\n",
      "Trainable params: 17,312,267\n",
      "Non-trainable params: 6\n",
      "_________________________________________________________________\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:35.292 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:24:35.293 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:24:35.936 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:24:35.937 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:24:35.999 | CRITICAL | m3tl.base_params:update_train_steps:454 - Updating train_steps to 1\n",
      "2021-06-17 13:24:36.554 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-17 13:24:36.555 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2702702702702703,\n",
      "    \"weibo_fake_multi_cls\": 0.2702702702702703,\n",
      "    \"weibo_masklm\": 0.1891891891891892,\n",
      "    \"weibo_premask_mlm\": 0.2702702702702703\n",
      "}\n",
      "2021-06-17 13:24:37.735 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:24:37.830 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:24:37.863 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:9 out of the last 9 calls to <function MultiModalBertModel.call at 0x7fd6ab4e2f80> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:39.238 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-17 13:24:39.239 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-17 13:24:39.240 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.iter\n",
      "WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_1\n",
      "WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_2\n",
      "WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.decay\n",
      "WARNING:tensorflow:A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:39.503 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:24:44.772 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - ETA: 0s - mean_acc: 3.5240 - weibo_fake_cls_acc: 0.3750 - weibo_fake_ner_acc: 0.2143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.4313 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.8266 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.8343 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0042 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9823"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:51.225 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1000 batches). You may need to use the repeat() function when building your dataset.\n",
      "1/1 [==============================] - 15s 15s/step - mean_acc: 3.5240 - weibo_fake_cls_acc: 0.3750 - weibo_fake_ner_acc: 0.2143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.4313 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.8266 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.8343 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0042 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9823 - val_loss: 17.1343 - val_mean_acc: 0.2500 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.5960 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_masklm/losses/0: 0.0000e+00 - val_BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9514\n",
      "Model: \"BertMultiTask\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "BertMultiTaskBody (BertMulti multiple                  4082696   \n",
      "_________________________________________________________________\n",
      "basic_mtl_5 (BasicMTL)       multiple                  0         \n",
      "_________________________________________________________________\n",
      "BertMultiTaskTop (BertMultiT multiple                  13229575  \n",
      "_________________________________________________________________\n",
      "sum_loss_combination_5 (SumL multiple                  0         \n",
      "=================================================================\n",
      "Total params: 17,312,273\n",
      "Trainable params: 17,312,267\n",
      "Non-trainable params: 6\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "\n",
    "# horovod train\n",
    "params.use_horovod = False\n",
    "_ = train_bert_multitask(\n",
    "    problem=problem,\n",
    "    num_epochs=1,\n",
    "    params=params,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict,\n",
    "    steps_per_epoch=1,\n",
    "    continue_training=False,\n",
    "    mirrored_strategy=False,\n",
    "    model_dir='./models/fresh_train'\n",
    ")\n",
    "\n",
    "model = train_bert_multitask(\n",
    "    problem=problem,\n",
    "    num_epochs=1,\n",
    "    params=params,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict,\n",
    "    steps_per_epoch=1,\n",
    "    continue_training=True,\n",
    "    mirrored_strategy=False,\n",
    "    run_eagerly=True\n",
    ")\n",
    "\n",
    "model = train_bert_multitask(\n",
    "    problem=problem,\n",
    "    num_epochs=1,\n",
    "    params=params,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict,\n",
    "    continue_training=True\n",
    ")\n",
    "\n",
    "# fresh train\n",
    "_ = train_bert_multitask(\n",
    "    problem=problem,\n",
    "    num_epochs=1,\n",
    "    params=params,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict,\n",
    "    steps_per_epoch=1,\n",
    "    continue_training=False,\n",
    "    mirrored_strategy=False,\n",
    "    model_dir='./models/fresh_train'\n",
    ")\n",
    "\n",
    "# transfer train\n",
    "params.init_checkpoint = './models/fresh_train'\n",
    "_ = train_bert_multitask(\n",
    "    problem=problem,\n",
    "    num_epochs=1,\n",
    "    params=params,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict,\n",
    "    steps_per_epoch=1,\n",
    "    continue_training=False,\n",
    "    mirrored_strategy=False\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving model for prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "def create_tensorspec_from_shape_type(infered_shape_and_type: Tuple[Dict[str, list], Dict[str, tf.dtypes.DType]]) -> Dict[str, tf.TensorSpec]:\n",
    "    \n",
    "    shape_dict, type_dict = infered_shape_and_type\n",
    "    tensorspec_dict = {}\n",
    "    for name in shape_dict.keys():\n",
    "        tensorspec_dict[name] = tf.TensorSpec(shape=[None for _ in shape_dict[name]], dtype=type_dict[name], name=name)\n",
    "    return tensorspec_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'text_input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_input_ids'), 'text_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_mask'), 'text_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_segment_ids'), 'image_input_ids': TensorSpec(shape=(None, None, None), dtype=tf.float32, name='image_input_ids'), 'image_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='image_mask'), 'image_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='image_segment_ids'), 'class_input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='class_input_ids'), 'class_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='class_mask'), 'class_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='class_segment_ids')}\n"
     ]
    }
   ],
   "source": [
    "test_tup = ({\n",
    "    'text_input_ids': [None, 3],\n",
    "    'text_mask': [None, 3],\n",
    "    'text_segment_ids': [None, 3],\n",
    "    'image_input_ids': [None, 5, 10],\n",
    "    'image_mask': [None, 5],\n",
    "    'image_segment_ids': [None, 5],\n",
    "    'class_input_ids': [None, 1],\n",
    "    'class_mask': [None, 1],\n",
    "    'class_segment_ids': [None, 1]},\n",
    "    {\n",
    "    'text_input_ids': tf.int32,\n",
    "    'text_mask': tf.int32,\n",
    "    'text_segment_ids': tf.int32,\n",
    "    'image_input_ids': tf.float32,\n",
    "    'image_mask': tf.int32,\n",
    "    'image_segment_ids': tf.int32,\n",
    "    'class_input_ids': tf.int32,\n",
    "    'class_mask': tf.int32,\n",
    "    'class_segment_ids': tf.int32})\n",
    "\n",
    "\n",
    "print(create_tensorspec_from_shape_type(test_tup))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@logger.catch\n",
    "def trim_checkpoint_for_prediction(problem: str,\n",
    "                                   input_dir: str,\n",
    "                                   output_dir: str,\n",
    "                                   problem_type_dict: Dict[str, str] = None,\n",
    "                                   overwrite=True,\n",
    "                                   fake_input_list=None,\n",
    "                                   params=None,\n",
    "                                   save_weights_only=True):\n",
    "\n",
    "    if overwrite and os.path.exists(output_dir):\n",
    "        rmtree(output_dir)\n",
    "    copytree(input_dir, output_dir, ignore=ignore_patterns(\n",
    "        'checkpoint', '*.index', '*.data-000*'))\n",
    "    base_dir, dir_name = os.path.split(output_dir)\n",
    "    if params is None:\n",
    "        params = Params()\n",
    "\n",
    "    if problem_type_dict:\n",
    "        params.register_multiple_problems(problem_type_dict=problem_type_dict)\n",
    "    params.from_json(os.path.join(input_dir, 'params.json'))\n",
    "    params.assign_problem(problem, base_dir=base_dir,\n",
    "                          dir_name=dir_name, predicting=True)\n",
    "\n",
    "    model = BertMultiTask(params)\n",
    "    if fake_input_list is None:\n",
    "        dummy_dataset = predict_input_fn(['fake']*5, params)\n",
    "    else:\n",
    "        dummy_dataset = predict_input_fn(fake_input_list*5, params)\n",
    "\n",
    "    batch_fake_data = next(dummy_dataset.as_numpy_iterator())\n",
    "    shape_type_dict = infer_shape_and_type_from_dict(batch_fake_data)\n",
    "    spec_dict = create_tensorspec_from_shape_type(infered_shape_and_type=shape_type_dict)\n",
    "\n",
    "    # _ = model(batch_fake_data,\n",
    "    #           mode=tf.estimator.ModeKeys.PREDICT)\n",
    "    # monkeypatch predict_step\n",
    "    # model.call = MethodType(tf.function(lambda x, y: model.call(x, mode=PREDICT)), model)\n",
    "    set_phase(PREDICT)\n",
    "    _ = model(batch_fake_data)\n",
    "    model.load_weights(os.path.join(input_dir, 'model'))\n",
    "    logger.critical(\"serving input sigantures: {}\".format(spec_dict))\n",
    "    if save_weights_only:\n",
    "        model.save_weights(os.path.join(params.ckpt_dir, 'model'))\n",
    "    else:\n",
    "        class ServingModule(tf.Module):\n",
    "            def __init__(self):\n",
    "                super(ServingModule, self).__init__()\n",
    "                self.model = model\n",
    "\n",
    "            @tf.function\n",
    "            def serve(self, x):\n",
    "                return self.model.call(x)\n",
    "        serving_module = ServingModule()\n",
    "        _ = serving_module.serve(batch_fake_data)\n",
    "        signatures = dict(\n",
    "            serving_default=serving_module.serve.get_concrete_function(spec_dict)\n",
    "        )\n",
    "        tf.saved_model.save(serving_module, os.path.join(params.ckpt_dir, 'serving'), signatures=signatures)\n",
    "    params.to_json()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Minimize checkpoint size for prediction.\n",
    "\n",
    "Since the original checkpoint contains optimizer's variable,\n",
    "        for instance, if the use adam, the checkpoint size will\n",
    "        be three times of the size of model weights. This function\n",
    "        will remove those unused variables in prediction to save space.\n",
    "\n",
    "Note: if the model is a multimodal model, you have to provide fake_input_list that\n",
    "        mimic the structure of real input. Otherwise modal embeddings will be randomly initialized.\n",
    "\n",
    "Args:\n",
    "- problem (str): problem\n",
    "- input_dir (str): input dir\n",
    "- output_dir (str): output dir\n",
    "- problem_type_dict (Dict[str, str], optional): problem type dict. Defaults to None.\n",
    "- fake_input_list (List, optional): fake input list to create dummy dataset\n",
    "- params (Params, optional): params "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 13:24:55.159 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:24:55.159 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:24:55.160 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:24:55.160 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:24:55.161 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:24:55.161 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:24:55.164 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:24:55.165 | WARNING  | m3tl.base_params:assign_problem:634 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-17 13:24:55.383 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:24:55.476 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:24:55.615 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:24:57.074 | CRITICAL | __main__:trim_checkpoint_for_prediction:43 - serving input sigantures: {'text_input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_input_ids'), 'text_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_mask'), 'text_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_segment_ids'), 'array_input_ids': TensorSpec(shape=(None, None, None), dtype=tf.float32, name='array_input_ids'), 'array_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='array_mask'), 'array_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='array_segment_ids'), 'cate_input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='cate_input_ids'), 'cate_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='cate_mask'), 'cate_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='cate_segment_ids')}\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "WARNING:absl:Found untraced functions such as embeddings_layer_call_fn, embeddings_layer_call_and_return_conditional_losses, position_embeddings_layer_call_fn, position_embeddings_layer_call_and_return_conditional_losses, token_type_embeddings_layer_call_fn while saving (showing 5 of 125). These functions will not be directly callable after loading.\n",
      "WARNING:absl:Found untraced functions such as embeddings_layer_call_fn, embeddings_layer_call_and_return_conditional_losses, position_embeddings_layer_call_fn, position_embeddings_layer_call_and_return_conditional_losses, token_type_embeddings_layer_call_fn while saving (showing 5 of 125). These functions will not be directly callable after loading.\n",
      "2021-06-17 13:25:39.080 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-17 13:25:39.080 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-17 13:25:39.081 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 13:25:39.082 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-17 13:25:39.082 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-17 13:25:39.082 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-17 13:25:39.083 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-17 13:25:39.084 | WARNING  | m3tl.base_params:assign_problem:634 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-17 13:25:39.306 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-17 13:25:39.404 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-17 13:25:39.558 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-17 13:25:40.708 | CRITICAL | __main__:trim_checkpoint_for_prediction:43 - serving input sigantures: {'text_input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_input_ids'), 'text_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_mask'), 'text_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='text_segment_ids'), 'array_input_ids': TensorSpec(shape=(None, None, None), dtype=tf.float32, name='array_input_ids'), 'array_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='array_mask'), 'array_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='array_segment_ids'), 'cate_input_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='cate_input_ids'), 'cate_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='cate_mask'), 'cate_segment_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='cate_segment_ids')}\n"
     ]
    }
   ],
   "source": [
    "# fake inputs\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "import numpy as np\n",
    "from m3tl.predefined_problems.test_data import generate_fake_data\n",
    "fake_inputs = [v for v, _ in generate_fake_data(output_format='gen_dict_tuple')]\n",
    "\n",
    "# save as SavedModel pb\n",
    "trim_checkpoint_for_prediction(problem=model.params.problem_str, input_dir=model.params.ckpt_dir,\n",
    "    output_dir=model.params.ckpt_dir+'_pred',\n",
    "    problem_type_dict=problem_type_dict, overwrite=True, fake_input_list=fake_inputs, save_weights_only=False)\n",
    "\n",
    "trim_checkpoint_for_prediction(\n",
    "    problem=problem, input_dir=model.params.ckpt_dir,\n",
    "    output_dir=model.params.ckpt_dir+'_pred',\n",
    "    problem_type_dict=problem_type_dict, overwrite=True, fake_input_list=fake_inputs)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@logger.catch\n",
    "def eval_bert_multitask(\n",
    "        problem='weibo_ner',\n",
    "        num_gpus=1,\n",
    "        model_dir='',\n",
    "        params=None,\n",
    "        problem_type_dict=None,\n",
    "        processing_fn_dict=None,\n",
    "        model=None,\n",
    "        run_eagerly=False):\n",
    "\n",
    "    if not model_dir and params is not None:\n",
    "        model_dir = params.ckpt_dir\n",
    "    params = get_params_ready(problem, num_gpus, model_dir,\n",
    "                              params, problem_type_dict, processing_fn_dict,\n",
    "                              mode='predict', json_path=os.path.join(model_dir, 'params.json'))\n",
    "    eval_dataset = train_eval_input_fn(params, mode=EVAL)\n",
    "    one_batch_data = next(eval_dataset.as_numpy_iterator())\n",
    "    eval_dataset = train_eval_input_fn(params, mode=EVAL)\n",
    "    mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "    if model is None:\n",
    "        model = create_keras_model(\n",
    "            mirrored_strategy=mirrored_strategy, params=params,\n",
    "            mode='eval', inputs_to_build_model=one_batch_data, run_eagerly=run_eagerly)\n",
    "    eval_dict = model.evaluate(eval_dataset, return_dict=True)\n",
    "    return eval_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate Multi-task Bert model\n",
    "\n",
    "Keyword Arguments:\n",
    "\n",
    "- problem (str, optional): problems to evaluate. Defaults to 'weibo_ner'.\n",
    "- num_gpus (int, optional): number of gpu to use. Defaults to 1.\n",
    "- model_dir (str, optional): model dir. Defaults to ''.\n",
    "- params (Params, optional): params. Defaults to None.\n",
    "- problem_type_dict (dict, optional): Key: problem name, value: problem type. Defaults to None.\n",
    "- processing_fn_dict (dict, optional): Key: problem name, value: problem data preprocessing fn. Defaults to None.\n",
    "- model (tf.keras.Model, optional): If not provided, it will be created with `create_keras_model`. Defaults to None.\n",
    "- run_eagerly (bool, optional): Whether to run model eagerly. Defaults to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:28:35.752 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 20:28:35.753 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-15 20:28:35.754 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 20:28:35.754 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 20:28:35.755 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 20:28:35.755 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-15 20:28:35.755 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 20:28:36.219 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-15 20:28:36.220 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2564102564102564,\n",
      "    \"weibo_fake_multi_cls\": 0.2564102564102564,\n",
      "    \"weibo_masklm\": 0.23076923076923078,\n",
      "    \"weibo_premask_mlm\": 0.2564102564102564\n",
      "}\n",
      "2021-06-15 20:28:37.241 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-15 20:28:37.243 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2564102564102564,\n",
      "    \"weibo_fake_multi_cls\": 0.2564102564102564,\n",
      "    \"weibo_masklm\": 0.23076923076923078,\n",
      "    \"weibo_premask_mlm\": 0.2564102564102564\n",
      "}\n",
      "2021-06-15 20:28:37.597 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-15 20:28:37.691 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-15 20:28:37.723 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-15 20:28:39.219 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-15 20:28:39.220 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-15 20:28:39.220 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "2021-06-15 20:28:40.317 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3/3 [==============================] - 5s 21ms/step - loss: 22.2266 - mean_acc: 0.5357 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.5714 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.2471 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.5657 - BertMultiTaskTop/weibo_fake_ner/losses/0: 0.3361 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0024 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9237\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:28:44.361 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 20:28:44.362 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-15 20:28:44.363 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 20:28:44.363 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 20:28:44.363 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 20:28:44.364 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-15 20:28:44.364 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 20:28:44.805 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-15 20:28:44.806 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2564102564102564,\n",
      "    \"weibo_fake_multi_cls\": 0.2564102564102564,\n",
      "    \"weibo_masklm\": 0.23076923076923078,\n",
      "    \"weibo_premask_mlm\": 0.2564102564102564\n",
      "}\n",
      "2021-06-15 20:28:45.832 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-15 20:28:45.833 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2564102564102564,\n",
      "    \"weibo_fake_multi_cls\": 0.2564102564102564,\n",
      "    \"weibo_masklm\": 0.23076923076923078,\n",
      "    \"weibo_premask_mlm\": 0.2564102564102564\n",
      "}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3/3 [==============================] - 1s 21ms/step - loss: 23.1741 - mean_acc: 0.4241 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.2857 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.4929 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.2098 - BertMultiTaskTop/weibo_fake_ner/losses/0: 1.7419 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0072 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.7933\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'loss': 23.17413330078125,\n",
       " 'mean_acc': 0.4241071343421936,\n",
       " 'weibo_fake_cls_acc': 0.5,\n",
       " 'weibo_fake_ner_acc': 0.2857142984867096,\n",
       " 'BertMultiTaskTop/weibo_fake_cls/losses/0': 1.8408839702606201,\n",
       " 'BertMultiTaskTop/weibo_fake_multi_cls/losses/0': 0.0,\n",
       " 'BertMultiTaskTop/weibo_fake_ner/losses/0': 1.7550007104873657,\n",
       " 'BertMultiTaskTop/weibo_masklm/losses/0': 10.002845764160156,\n",
       " 'BertMultiTaskTop/weibo_premask_mlm/losses/0': 9.788535118103027}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# remove original ckpt path to make sure model can be init from a different path\n",
    "import shutil\n",
    "shutil.rmtree(model.params.ckpt_dir)\n",
    "\n",
    "eval_bert_multitask(problem=problem, params=params,\n",
    "                    problem_type_dict=problem_type_dict, processing_fn_dict=processing_fn_dict,\n",
    "                    model_dir=model.params.ckpt_dir+'_pred')\n",
    "\n",
    "# provide model instead of dir\n",
    "eval_bert_multitask(problem=problem, params=params,\n",
    "                    problem_type_dict=problem_type_dict, processing_fn_dict=processing_fn_dict,\n",
    "                    model=model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:28:47.317 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 20:28:47.318 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-15 20:28:47.318 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 20:28:47.319 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 20:28:47.319 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 20:28:47.319 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-15 20:28:47.320 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 20:28:47.768 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-15 20:28:47.769 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2564102564102564,\n",
      "    \"weibo_fake_multi_cls\": 0.2564102564102564,\n",
      "    \"weibo_masklm\": 0.23076923076923078,\n",
      "    \"weibo_premask_mlm\": 0.2564102564102564\n",
      "}\n",
      "2021-06-15 20:28:48.629 | INFO     | m3tl.input_fn:train_eval_input_fn:59 - sampling weights: \n",
      "2021-06-15 20:28:48.629 | INFO     | m3tl.input_fn:train_eval_input_fn:60 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.2564102564102564,\n",
      "    \"weibo_fake_multi_cls\": 0.2564102564102564,\n",
      "    \"weibo_masklm\": 0.23076923076923078,\n",
      "    \"weibo_premask_mlm\": 0.2564102564102564\n",
      "}\n",
      "2021-06-15 20:28:48.909 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-15 20:28:49.003 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "2021-06-15 20:28:49.035 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-15 20:28:50.458 | CRITICAL | m3tl.model_fn:compile:271 - Initial lr: 0.0\n",
      "2021-06-15 20:28:50.459 | CRITICAL | m3tl.model_fn:compile:272 - Train steps: 1\n",
      "2021-06-15 20:28:50.459 | CRITICAL | m3tl.model_fn:compile:273 - Warmup steps: 0\n",
      "2021-06-15 20:28:51.586 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      1/Unknown - 2s 2s/step - loss: 23.5573 - mean_acc: 0.3690 - weibo_fake_cls_acc: 0.1667 - weibo_fake_ner_acc: 0.5714 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9506 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.7102 - BertMultiTaskTop/weibo_fake_ner/losses/0: 0.9680 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0068 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9217"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:28:52.575 | INFO     | m3tl.utils:set_phase:478 - Setting phase to eval\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - 3s 876ms/step - loss: 19.5134 - mean_acc: 0.4524 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.5714 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.6026 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 0.5701 - BertMultiTaskTop/weibo_fake_ner/losses/0: 0.9949 - BertMultiTaskTop/weibo_masklm/losses/0: 9.9947 - BertMultiTaskTop/weibo_premask_mlm/losses/0: 9.9217\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'loss': 17.491474151611328,\n",
       " 'mean_acc': 0.4523809552192688,\n",
       " 'weibo_fake_cls_acc': 0.5,\n",
       " 'weibo_fake_ner_acc': 0.5714285969734192,\n",
       " 'BertMultiTaskTop/weibo_fake_cls/losses/0': 0.4286690652370453,\n",
       " 'BertMultiTaskTop/weibo_fake_multi_cls/losses/0': 0.0,\n",
       " 'BertMultiTaskTop/weibo_fake_ner/losses/0': 1.0083808898925781,\n",
       " 'BertMultiTaskTop/weibo_masklm/losses/0': 9.988622665405273}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "# eager mode test\n",
    "eval_bert_multitask(problem=problem, params=params,\n",
    "                    problem_type_dict=problem_type_dict, processing_fn_dict=processing_fn_dict,\n",
    "                    model_dir=model.params.ckpt_dir, run_eagerly=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exprort\n",
    "\n",
    "\n",
    "def arr_to_str(inp_arr: np.ndarray) -> str:\n",
    "    l = inp_arr.tolist()\n",
    "    l = [json.dumps(f) for f in l]\n",
    "    return l\n",
    "\n",
    "\n",
    "def decode_predictions(pred: Dict[str, np.ndarray], params: Params, array_as_str=False) -> Dict[str, Union[int, float, np.ndarray, list, str]]:\n",
    "    parsed_pred = dict()\n",
    "    problem_list = params.problem_list\n",
    "    label_encoder_dict = {p: get_or_make_label_encoder(\n",
    "        params=params, problem=p, mode=PREDICT) for p in problem_list}\n",
    "    for problem, problem_pred_array in pred.items():\n",
    "\n",
    "        # addtional outputs\n",
    "        if problem not in problem_list:\n",
    "            if isinstance(problem_pred_array, np.ndarray):\n",
    "                if array_as_str:\n",
    "                    parsed_pred[problem] = arr_to_str(problem_pred_array)\n",
    "                else:\n",
    "                    parsed_pred[problem] = problem_pred_array\n",
    "            else:\n",
    "                parsed_pred[problem] = problem_pred_array\n",
    "            continue\n",
    "\n",
    "        label_encoder = label_encoder_dict[problem]\n",
    "\n",
    "        support_problem_type = [\n",
    "            'multi_cls',\n",
    "            'cls',\n",
    "            'seq_tag',\n",
    "            'regression',\n",
    "            'masklm',\n",
    "            'premask_mlm',\n",
    "            'vectorfit'\n",
    "        ]\n",
    "\n",
    "        problem_type = params.get_problem_type(problem=problem)\n",
    "        if problem_type not in support_problem_type:\n",
    "            logger.warning(\"trying to decode prediction of unsupported problem type\"\n",
    "            \" {}, if any error raised, please disable decode prediction.\".format(problem_type))\n",
    "\n",
    "        is_multi_cls = params.get_problem_type(problem=problem) == 'multi_cls'\n",
    "        is_cls = params.get_problem_type(problem=problem) == 'cls'\n",
    "        is_seq_tag = params.get_problem_type(problem=problem) == 'seq_tag'\n",
    "        is_regression = params.get_problem_type(\n",
    "            problem=problem) == 'regression'\n",
    "\n",
    "        if is_regression:\n",
    "            parsed_pred[problem] = problem_pred_array\n",
    "            continue\n",
    "\n",
    "        # get pred from prob\n",
    "        if is_multi_cls:\n",
    "            problem_pred = problem_pred_array >= 0.5\n",
    "        elif is_cls or is_seq_tag:\n",
    "            problem_pred = np.argmax(problem_pred_array, axis=-1)\n",
    "            # problem_pred = problem_pred_array\n",
    "        else:\n",
    "            problem_pred = problem_pred_array\n",
    "\n",
    "        # sequence labels\n",
    "        if is_seq_tag:\n",
    "            parsed_problem_pred = np.apply_along_axis(\n",
    "                label_encoder.inverse_transform, axis=1, arr=problem_pred)\n",
    "        else:\n",
    "            if isinstance(label_encoder, MultiLabelBinarizer) or isinstance(label_encoder, LabelEncoder):\n",
    "                parsed_problem_pred = label_encoder.inverse_transform(\n",
    "                    problem_pred)\n",
    "            elif isinstance(label_encoder, PreTrainedTokenizer):\n",
    "                parsed_problem_pred = np.apply_along_axis(\n",
    "                    label_encoder.convert_ids_to_tokens, axis=1, arr=problem_pred\n",
    "                )\n",
    "            else:\n",
    "                parsed_problem_pred = problem_pred_array\n",
    "\n",
    "        parsed_pred[problem] = parsed_problem_pred\n",
    "    return parsed_pred\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@logger.catch\n",
    "def predict_bert_multitask(\n",
    "        inputs,\n",
    "        problem='weibo_ner',\n",
    "        model_dir='',\n",
    "        params: Params = None,\n",
    "        problem_type_dict: Dict[str, str] = None,\n",
    "        processing_fn_dict: Dict[str, Callable] = None,\n",
    "        model: tf.keras.Model = None,\n",
    "        return_model=False,\n",
    "        run_eagerly=False,\n",
    "        mirrored_strategy=None,\n",
    "        decode_prediction=False):\n",
    "    \"\"\"Use Multi-task Bert model to do prediction\n",
    "\n",
    "        Args:\n",
    "        - inputs (Iterable): Iterable of inputs\n",
    "        - problem (str, optional): problems to predict. Defaults to 'weibo_ner'.\n",
    "        - model_dir (str, optional): model dir. Defaults to ''.\n",
    "        - params (Params, optional): params. Defaults to None.\n",
    "        - problem_type_dict (Dict[str, str], optional): Key: problem name, value: problem type.. Defaults to None.\n",
    "        - processing_fn_dict (Dict[str, Callable], optional): Key: problem name, value: problem data preprocessing fn. Defaults to None.\n",
    "        - model (tf.keras.Model, optional): If not provided, it will be created with `create_keras_model`. Defaults to None.\n",
    "        - return_model (bool, optional): Whether return model, if True, function will return (pred, model) tuple. Defaults to False.\n",
    "        - run_eagerly (bool, optional): Whether to run model eagerly. Defaults to False.\n",
    "        - mirrored_strategy (optional): mirrored strategy for distribute prediction. Defaults to None.\n",
    "        - decode_prediction (bool, optional): whether to decode predictions. Defaults to False.\n",
    "    \"\"\"\n",
    "    set_phase(PREDICT)\n",
    "    if params is None:\n",
    "        params = Params()\n",
    "    if not model_dir and params is not None:\n",
    "        model_dir = params.ckpt_dir\n",
    "    params = get_params_ready(problem, 1, model_dir,\n",
    "                              params, problem_type_dict, processing_fn_dict,\n",
    "                              mode='predict', json_path=os.path.join(model_dir, 'params.json'))\n",
    "\n",
    "    logger.info('Checkpoint dir: {}'.format(params.ckpt_dir))\n",
    "    time.sleep(3)\n",
    "\n",
    "    pred_dataset = predict_input_fn(inputs, params)\n",
    "    one_batch_data = next(pred_dataset.as_numpy_iterator())\n",
    "    pred_dataset = predict_input_fn(inputs, params)\n",
    "\n",
    "    if model is None:\n",
    "        model = create_keras_model(\n",
    "            mirrored_strategy=mirrored_strategy, params=params,\n",
    "            mode='predict', inputs_to_build_model=one_batch_data,\n",
    "            run_eagerly=run_eagerly)\n",
    "\n",
    "    if mirrored_strategy is not None:\n",
    "        with mirrored_strategy.scope():\n",
    "            pred = model.predict(pred_dataset)\n",
    "    else:\n",
    "        pred = model.predict(pred_dataset)\n",
    "\n",
    "    if decode_prediction:\n",
    "        pred = decode_predictions(pred=pred, params=model.params)\n",
    "\n",
    "    if return_model:\n",
    "        return pred, model\n",
    "    return pred\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:28:53.975 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-15 20:28:53.976 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 20:28:53.977 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-15 20:28:53.977 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 20:28:53.978 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 20:28:53.978 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 20:28:53.978 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-15 20:28:53.979 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 20:28:54.017 | INFO     | __main__:predict_bert_multitask:39 - Checkpoint dir: models/weibo_fake_cls_weibo_fake_multi_cls_weibo_fake_ner_weibo_masklm_weibo_premask_mlm_ckpt_pred\n",
      "2021-06-15 20:28:59.407 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-15 20:28:59.474 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-15 20:29:00.488 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "pred, model = predict_bert_multitask(\n",
    "    problem='weibo_fake_ner',\n",
    "    inputs=fake_inputs*20, model_dir=model.params.ckpt_dir,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict, return_model=True,\n",
    "    params=params)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:29:02.664 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-15 20:29:02.665 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 20:29:02.666 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-15 20:29:02.666 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 20:29:02.666 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 20:29:02.667 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 20:29:02.667 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-15 20:29:02.668 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 20:29:02.707 | INFO     | __main__:predict_bert_multitask:39 - Checkpoint dir: models/weibo_fake_cls_weibo_fake_multi_cls_weibo_fake_ner_weibo_masklm_weibo_premask_mlm_ckpt_pred\n",
      "2021-06-15 20:29:06.155 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-15 20:29:06.222 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-15 20:29:07.146 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "# eager mode test\n",
    "pred, model = predict_bert_multitask(\n",
    "    problem='weibo_fake_ner',\n",
    "    inputs=fake_inputs*20, model_dir=model.params.ckpt_dir,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict, return_model=True,\n",
    "    params=params, run_eagerly=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 20:29:09.212 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "2021-06-15 20:29:09.213 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 20:29:09.213 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_cws, problem type: seq_tag\n",
      "2021-06-15 20:29:09.214 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 20:29:09.214 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 20:29:09.215 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 20:29:09.215 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_pretrain, problem type: pretrain\n",
      "2021-06-15 20:29:09.216 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 20:29:09.255 | INFO     | __main__:predict_bert_multitask:39 - Checkpoint dir: models/weibo_fake_cls_weibo_fake_multi_cls_weibo_fake_ner_weibo_masklm_weibo_premask_mlm_ckpt_pred\n",
      "2021-06-15 20:29:12.688 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-15 20:29:12.791 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "2021-06-15 20:29:13.769 | INFO     | m3tl.utils:set_phase:478 - Setting phase to infer\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "pred, model = predict_bert_multitask(\n",
    "    problem='weibo_fake_ner|weibo_fake_cls|weibo_fake_multi_cls|weibo_premask_mlm',\n",
    "    inputs=fake_inputs*20, model_dir=model.params.ckpt_dir,\n",
    "    problem_type_dict=problem_type_dict,\n",
    "    processing_fn_dict=processing_fn_dict, return_model=True,\n",
    "    params=params,\n",
    "    decode_prediction=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
